from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.models import UserMessage
from autogen_agentchat.messages import MultiModalMessage
from autogen_core import Image as AGImage
from autogen_agentchat.base import TaskResult
from autogen_agentchat.conditions import ExternalTermination, TextMentionTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_core import CancellationToken
import PIL
from autogen_core import Image
from io import BytesIO
import asyncio
import requests
from typing import Literal
from pydantic import BaseModel


# Define a model client. You can use other model client that implements
# the `ChatCompletionClient` interface.


model_client = OpenAIChatCompletionClient(
    model="qwen-vl-max",
    api_key="sk-925b8bbb82424b74a8de940d2dc5a6ce",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model_info={
        "vision": True,
        "function_calling": True,
        "json_output": False,
        "family": "unknown",
        "structured_output": True,
        "max_tokens": 8192,
        "top_p":0.8
    }
    # api_key="YOUR_API_KEY",
)

primary_agent = AssistantAgent(
    "primary",
    model_client=model_client,
    system_message="You are a helpful AI assistant.",
)

# Create the critic agent.
critic_agent = AssistantAgent(
    "critic",
    model_client=model_client,
    system_message="Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.",
)

# Define a termination condition that stops the task if the critic approves.
text_termination = TextMentionTermination("APPROVE")

external_termination = ExternalTermination()

# Create a team with the primary and critic agents.
team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=text_termination | external_termination)

async def main():
    #result = await team.run(task="写一首关于秋天的诗")
    #print(result)
    
    #await Console(team.run_stream(task="写一首关于秋天的诗"))

    run = asyncio.create_task(Console(team.run_stream(task="Write a short poem about the fall season.")))

    # Wait for some time.
    #await asyncio.sleep(0.1)

    # Stop the team.
    #external_termination.set()

    # Wait for the team to finish.
    await run

    await Console(team.run_stream(task="将这首诗用中文唐诗风格写一遍。"))  # Resume the team to continue the last task.


asyncio.run(main())

